from trans import TransBelt
from ultralytics import YOLO
from qt import UI
import time
import cv2

trans = TransBelt()             # 初始化传送带
yolo = YOLO('./model/yolo11n.onnx')    # 初始化模型
qt = UI()                       # 初始化QT界面
cam = cv2.VideoCapture(0)


def wait_garbage_in() -> None:
    while (True):
        flag, img = cam.read()
        if not flag:
            continue
        res = yolo.predict(img)
        if len(res[0].summary()) > 0:
            break


def check_volume() -> bool:
    return False


if __name__ == "__main__":
    # 主循环
    while (True):
        wait_garbage_in()           # 阻塞等待垃圾投入
        qt.working()
        while (True):
            flag, img = cam.read()
            if not flag:
                continue
            items = yolo.predict(img)
            if (len(items[0].summary()) == 0):
                break
            # 检查是否有掉落 没有则返回-1 有则返回垃圾种类编号
            tmp = trans.is_there_anything_dropped()
            if (tmp != -1):
                trans.wait_post_process()             # 掉落后处理
                qt.one_garbage_in(tmp)
                # garbages.append(tmp)
            trans.run(items[0])
        qt.summary()
        time.sleep(10)
        qt.idle()
        volume = check_volume()
        if volume:
            qt.set_full()
            break
    while (True):
        continue
